package org.qq.util;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

public class SQLLineageMerger {

    public static String mergeSQLLineage(List<String> lineageStrings, String columnName) {
        if (lineageStrings.isEmpty()) {
            return null;
        }

        String result = mergeSQLLineage(lineageStrings);
        if (columnName == null) return result;

        ObjectMapper objectMapper = new ObjectMapper();
        try {
            JsonNode rootNode = objectMapper.readTree(result);
            JsonNode dataNode = rootNode.path("data");
            JsonNode withProcessDataNode = dataNode.path("withProcessData");
            ArrayNode dataArrayNode = (ArrayNode) withProcessDataNode.path("data");

            List<String> relevantLineage = new ArrayList<>();
            for (JsonNode lineageNode : dataArrayNode) {
                JsonNode targetFieldNode = lineageNode.path("targetField");
                String targetFieldName = targetFieldNode.path("fieldName").asText();

                if (targetFieldName.endsWith("." + columnName)) {
                    relevantLineage.add(targetFieldName);
                    continue;
                }

                ArrayNode refFieldsArray = (ArrayNode) lineageNode.path("refFields");
                for (JsonNode refFieldNode : refFieldsArray) {
                    String refFieldName = refFieldNode.path("fieldName").asText();
                    if (refFieldName.endsWith("." + columnName)) {
                        relevantLineage.add(refFieldName);
                        break;
                    }
                }
            }

            if (relevantLineage.isEmpty()) {
                return null;
            }

            return filterLineageData(result, relevantLineage, columnName, objectMapper);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    private static String filterLineageData(String originalResult, List<String> relevantFields, String columnName, ObjectMapper objectMapper) {
        try {
            JsonNode rootNode = objectMapper.readTree(originalResult);
            JsonNode dataNode = rootNode.path("data");
            JsonNode withProcessDataNode = dataNode.path("withProcessData");

            if (withProcessDataNode.isObject()) {
                ObjectNode withObjectNode = (ObjectNode) withProcessDataNode;
                ArrayNode dataArrayNode = (ArrayNode) withObjectNode.path("data");
                ArrayNode filteredDataArray = objectMapper.createArrayNode();

                for (JsonNode lineageNode : dataArrayNode) {
                    JsonNode targetFieldNode = lineageNode.path("targetField");
                    String targetFieldName = targetFieldNode.path("fieldName").asText();

                    if (targetFieldName.endsWith("." + columnName)) {
                        filteredDataArray.add(lineageNode);
                        continue;
                    }

                    ArrayNode refFieldsArray = (ArrayNode) lineageNode.path("refFields");
                    boolean hasRelevantRefField = false;
                    for (JsonNode refFieldNode : refFieldsArray) {
                        String refFieldName = refFieldNode.path("fieldName").asText();
                        if (refFieldName.endsWith("." + columnName)) {
                            hasRelevantRefField = true;
                            break;
                        }
                    }

                    if (hasRelevantRefField) {
                        filteredDataArray.add(lineageNode);
                    }
                }

                withObjectNode.set("data", filteredDataArray);
                withObjectNode.put("size", filteredDataArray.size());
            }

            return objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(rootNode);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String mergeSQLLineage(List<String> lineageStrings) {
        ObjectMapper objectMapper = new ObjectMapper();
        Map<String, List<String>> lineageMap = new HashMap<>();
        Set<String> allFields = new HashSet<>();
        Set<String> targetFields = new HashSet<>();
        Map<String, Integer> levelMap = new ConcurrentHashMap<>();

        parseLineageStrings(lineageStrings, objectMapper, lineageMap, allFields, targetFields);
        levelMap = calculateLevels(allFields, lineageMap);

        Map<Integer, Map<String, List<String>>> levelToTableFields = groupFieldsByLevelAndTable(allFields, levelMap);
        Map<Integer, Map<String, Integer>> tableIndexMap = assignTableIndexes(levelToTableFields);

        return buildResult(objectMapper, lineageMap, levelMap, tableIndexMap, targetFields, levelToTableFields);
    }

    private static void parseLineageStrings(List<String> lineageStrings, ObjectMapper objectMapper,
                                            Map<String, List<String>> lineageMap, Set<String> allFields, Set<String> targetFields) {
        for (String lineageString : lineageStrings) {
            try {
                JsonNode rootNode = objectMapper.readTree(lineageString);
                String schema = rootNode.path("schema").asText();
                String table = rootNode.path("table").asText();
                String fullTableName = schema + "." + table;

                JsonNode columnsNode = rootNode.path("columns");
                for (JsonNode columnNode : columnsNode) {
                    String columnName = columnNode.path("column").asText();
                    JsonNode sourceColumnsNode = columnNode.path("sourceColumns");

                    String targetField = fullTableName + "." + columnName;
                    targetFields.add(targetField);
                    allFields.add(targetField);

                    if (sourceColumnsNode.size() > 0) {
                        List<String> sourceFieldsList = new ArrayList<>();
                        for (JsonNode sourceColumnNode : sourceColumnsNode) {
                            String sourceTableName = sourceColumnNode.path("tableName").asText();
                            String sourceColumnName = sourceColumnNode.path("columnName").asText();
                            String sourceField = sourceTableName + "." + sourceColumnName;
                            sourceFieldsList.add(sourceField);
                            allFields.add(sourceField);
                        }
                        lineageMap.put(targetField, sourceFieldsList);
                    }
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    private static Map<String, Integer> calculateLevels(Set<String> allFields, Map<String, List<String>> lineageMap) {
        Map<String, Integer> levelMap = new ConcurrentHashMap<>();
        allFields.forEach(field -> levelMap.put(field, 0));

        Queue<String> queue = new LinkedList<>();
        Set<String> visited = new HashSet<>();

        lineageMap.keySet().forEach(field -> {
            queue.add(field);
            visited.add(field);
        });

        while (!queue.isEmpty()) {
            String currentField = queue.poll();
            int currentLevel = levelMap.get(currentField);

            if (lineageMap.containsKey(currentField)) {
                for (String refField : lineageMap.get(currentField)) {
                    if (levelMap.get(refField) < currentLevel + 1) {
                        levelMap.put(refField, currentLevel + 1);
                        if (!visited.contains(refField)) {
                            queue.add(refField);
                            visited.add(refField);
                        }
                    }
                }
            }
        }
        return levelMap;
    }

    private static Map<Integer, Map<String, List<String>>> groupFieldsByLevelAndTable(Set<String> allFields, Map<String, Integer> levelMap) {
        Map<Integer, Map<String, List<String>>> levelToTableFields = new HashMap<>();
        for (String field : allFields) {
            int level = levelMap.get(field);
            String table = field.split("\\.", 2)[0];

            levelToTableFields
                    .computeIfAbsent(level, k -> new HashMap<>())
                    .computeIfAbsent(table, k -> new ArrayList<>())
                    .add(field);
        }
        return levelToTableFields;
    }

    private static Map<Integer, Map<String, Integer>> assignTableIndexes(Map<Integer, Map<String, List<String>>> levelToTableFields) {
        Map<Integer, Map<String, Integer>> tableIndexMap = new HashMap<>();
        for (Map.Entry<Integer, Map<String, List<String>>> levelEntry : levelToTableFields.entrySet()) {
            int level = levelEntry.getKey();
            Map<String, List<String>> tables = levelEntry.getValue();

            List<String> sortedTables = new ArrayList<>(tables.keySet());
            Collections.sort(sortedTables);

            Map<String, Integer> indexMap = new HashMap<>();
            for (int i = 0; i < sortedTables.size(); i++) {
                indexMap.put(sortedTables.get(i), i);
            }
            tableIndexMap.put(level, indexMap);
        }
        return tableIndexMap;
    }

    private static String buildResult(ObjectMapper mapper,
                                      Map<String, List<String>> lineageMap,
                                      Map<String, Integer> levelMap,
                                      Map<Integer, Map<String, Integer>> tableIndexMap,
                                      Set<String> targetFields,
                                      Map<Integer, Map<String, List<String>>> levelToTableFields) {
        ObjectNode resultNode = mapper.createObjectNode();
        resultNode.put("code", 0);

        // Build withProcessData
        ObjectNode withProcessNode = buildWithProcessData(mapper, lineageMap, targetFields, levelMap, tableIndexMap);

        // Build noProcessData
        ObjectNode noProcessNode = buildNoProcessData(mapper, lineageMap, targetFields, levelMap, tableIndexMap);

        ObjectNode dataNode = mapper.createObjectNode();
        dataNode.set("withProcessData", withProcessNode);
        dataNode.set("noProcessData", noProcessNode);
        resultNode.set("data", dataNode);

        resultNode.put("errno", 0);
        resultNode.put("error", "");
        resultNode.put("request_id", UUID.randomUUID().toString().replace("-", ""));

        try {
            return mapper.writerWithDefaultPrettyPrinter().writeValueAsString(resultNode);
        } catch (IOException e) {
            e.printStackTrace();
            return "{}";
        }
    }

    private static ObjectNode buildWithProcessData(ObjectMapper mapper,
                                               Map<String, List<String>> lineageMap,
                                               Set<String> targetFields,
                                               Map<String, Integer> levelMap,
                                               Map<Integer, Map<String, Integer>> tableIndexMap) {
        ObjectNode withProcessNode = mapper.createObjectNode();
        ArrayNode dataArray = mapper.createArrayNode();

        for (String target : targetFields) {
            ObjectNode lineageNode = mapper.createObjectNode();
            ArrayNode refFieldsArray = mapper.createArrayNode();

            // Add target field
            ObjectNode targetNode = createFieldNode(mapper, target, levelMap, tableIndexMap);
            if (!lineageMap.containsKey(target)) {
                targetNode.put("final", true);
            }
            lineageNode.set("targetField", targetNode);

            // Add references if they exist
            if (lineageMap.containsKey(target)) {
                for (String ref : lineageMap.get(target)) {
                    ObjectNode refNode = createFieldNode(mapper, ref, levelMap, tableIndexMap);
                    refNode.put("final", !lineageMap.containsKey(ref));
                    refFieldsArray.add(refNode);
                }
            }

            lineageNode.set("refFields", refFieldsArray);
            dataArray.add(lineageNode);
        }

        withProcessNode.set("data", dataArray);
        withProcessNode.put("size", dataArray.size());
        withProcessNode.put("level", levelMap.values().stream().max(Integer::compare).orElse(0));
        return withProcessNode;
    }

    private static ObjectNode buildNoProcessData(ObjectMapper mapper,
                                             Map<String, List<String>> lineageMap,
                                             Set<String> targetFields,
                                             Map<String, Integer> levelMap,
                                             Map<Integer, Map<String, Integer>> tableIndexMap) {
        ObjectNode noProcessNode = mapper.createObjectNode();
        ArrayNode dataArray = mapper.createArrayNode();

        // 1. 识别最终目标字段（包括没有源字段的目标字段）
        Set<String> allSources = lineageMap.values().stream()
            .flatMap(List::stream)
            .collect(Collectors.toSet());

        Set<String> finalTargets = new HashSet<>(targetFields);
        finalTargets.removeAll(allSources); // 保留没有被其他字段引用的目标

        // 2. 为所有最终目标字段构建简化血缘
        for (String target : finalTargets) {
            ObjectNode lineageNode = mapper.createObjectNode();
            ArrayNode refFieldsArray = mapper.createArrayNode();

            // 目标字段
            ObjectNode targetNode = createFieldNode(mapper, target, levelMap, tableIndexMap);
            if (!lineageMap.containsKey(target)) {
                targetNode.put("final", true);
            }
            lineageNode.set("targetField", targetNode);

            // 如果有源字段，添加源字段
            if (lineageMap.containsKey(target)) {
                Set<String> roots = findRootSources(target, lineageMap);
                for (String root : roots) {
                    ObjectNode refNode = createFieldNode(mapper, root, levelMap, tableIndexMap);
                    refNode.put("final", true);
                    refFieldsArray.add(refNode);
                }
            }
            // 即使没有源字段，也添加到结果中
            lineageNode.set("refFields", refFieldsArray);
            dataArray.add(lineageNode);
        }

        noProcessNode.set("data", dataArray);
        noProcessNode.put("size", dataArray.size());
        noProcessNode.put("level", calculateMaxLevel(dataArray));
        return noProcessNode;
    }


    private static Set<String> findRootSources(String field, Map<String, List<String>> lineageMap) {
        Set<String> roots = new HashSet<>();
        Deque<String> stack = new ArrayDeque<>();
        Set<String> visited = new HashSet<>();

        stack.push(field);
        while (!stack.isEmpty()) {
            String current = stack.pop();
            if (!lineageMap.containsKey(current)) {
                roots.add(current);
                continue;
            }

            for (String source : lineageMap.get(current)) {
                if (visited.add(source)) {
                    stack.push(source);
                } else {
                    roots.add(source); // Handle cyclic dependencies
                }
            }
        }
        return roots;
    }

    private static ObjectNode createFieldNode(ObjectMapper mapper,
                                              String field,
                                              Map<String, Integer> levelMap,
                                              Map<Integer, Map<String, Integer>> tableIndexMap) {
        ObjectNode node = mapper.createObjectNode();
        node.put("fieldName", field);
        node.put("final", false);
        node.put("index", getTableIndex(field, tableIndexMap, levelMap));
        node.put("level", levelMap.getOrDefault(field, 0));
        return node;
    }

    private static int getTableIndex(String field,
                                     Map<Integer, Map<String, Integer>> tableIndexMap,
                                     Map<String, Integer> levelMap) {
        int level = levelMap.getOrDefault(field, 0);
        String table = field.split("\\.", 2)[0];
        return tableIndexMap.getOrDefault(level, Collections.emptyMap())
                .getOrDefault(table, 0);
    }

    private static int calculateMaxLevel(ArrayNode data) {
        int maxLevel = 0;
        for (JsonNode node : data) {
            int level = node.path("targetField").path("level").asInt();
            if (level > maxLevel) maxLevel = level;
        }
        return maxLevel;
    }
}